function [effect,effect_var,Z,llik_trace,elbo] = var_spike_slab_logit(Y,X,prior_lodds,prior_var,MAXEPOCH,effect,Z)
%
% [EFFECT,EFFECTVAR,Z,llik,elbo] = FUNC(Y,X,PRIOR_LODDS,PRIOR_VAR,MAXEPOCH,EFFECT,Z)
%
% Inference of association model
%
% Y(i,:) ~ logit(X(i,:)*EFFECT')
%
% Y           = sample x phenotype probabilistic/binary matrix
% X           = sample x feature data matrix (real number)
% PRIOR_LODDS = pheno x feature prior log-odds
% PRIOR_VAR   = prior variance
% MAXEPOCH    = maximum number of iterations
% EFFECT      = pheno x feature effect size
% EFFECTVAR   = pheno x feature effect size variance
% Z           = feature x 1 selection
% llik        = log-likelihood trace
% elbo        = expected log-likelihood bound
%
% code: Yongjin Park, ypp@csail.mit.edu
%

    TOL = 1e-4;
    P_TRUNC = 1e-5;

    % ================================================================

    [Nsample,Npheno] = size(Y);
    [nsample,Nfeat] = size(X);
    assert(Nsample == nsample);

    % initialize
    if nargin > 5,
        assert(Npheno == size(effect,1) && Nfeat == size(effect,2));
        assert(Nfeat == size(Z,1) && size(Z,2) == 1);
    else
        effect = zeros(Npheno,Nfeat,'single');
        Z = zeros(Nfeat,1,'single');
    end

    invtausq_effect = 1./prior_var;
    effect_var = zeros(Npheno,Nfeat,'single');
    effect_var_conditioned = zeros(Npheno,Nfeat,'single');
    effect_sq_conditioned = zeros(Npheno,Nfeat,'single');

    assert(all(size(prior_lodds) == size(Z)));

    llik_trace = NaN(MAXEPOCH,1,'single');

    for iter = 1:MAXEPOCH,

        % construct local quadratic approximation
        F = X * bsxfun(@times,effect',Z);
        P = 1./(1+exp(-F));
        P(P < P_TRUNC) = P_TRUNC;
        P(P > 1 - P_TRUNC) = 1 - P_TRUNC;

        W = P.*(1-P);
        workRes = F + (Y - P) ./ W;

        % coordinate-descent step
        for jj = 1:Nfeat,

            % [effect size of probes]
            denom_effect = invtausq_effect + W'*X(:,jj).^2;
            num_effect = (W.* (workRes - F + Z(jj)*X(:,jj)*effect(:,jj)'))'*X(:,jj);
            aa = num_effect ./ denom_effect;

            % E[effect^2|Z=1] * P(Z=1) - E[effect|Z=1]^2 * P(Z=1)^2
            effect_var_conditioned(:,jj) = 1./denom_effect;
            effect_var(:,jj) = (1./denom_effect + aa.^2)*Z(jj) - (aa.^2)*(Z(jj)^2);
            effect_sq_conditioned(:,jj) = 1./denom_effect + aa.^2;

            % update F
            F = F + Z(jj)*X(:,jj)*(aa - effect(:,jj))';
            effect(:,jj) = aa;

            % [probe selection]
            mass = X(:,jj)'*(W .* (workRes - F + Z(jj)*X(:,jj)*effect(:,jj)')) * effect(:,jj);
            mass = mass - 0.5 * (X(:,jj).^2)' * W * (effect(:,jj).^2);
            mass = mass + prior_lodds(jj);
            mass = mass + 0.5 * sum(log(W'*(X(:,jj).^2) ./ invtausq_effect + 1));

            zz = 1./(1+exp(-mass));

            % update F
            F = F + (zz - Z(jj))*X(:,jj)*effect(:,jj)';
            Z(jj) = zz;

        end

        llik = sum(sum(P.*log(P))) + sum(sum((1-P).*log(1-P)));
        llik_trace(iter) = llik;

        if iter > 10,
            llik_prev = mean(llik_trace((iter-10):(iter-6)));
            llik_curr = mean(llik_trace((iter-5):(iter-1)));

            if abs(llik_prev - llik_curr)/abs(TOL + llik_curr) < TOL,
                llik_trace = llik_trace(1:iter);
                fprintf(2,'Converged\r');
                break;
            end
        end

        fprintf(2,'Iter = %03d, LLIK = %.4e\r',iter, llik);

    end

    % here we assume quadratic from, which we used in variational
    % inference -- also comparable with gausisan model 
    F = X * bsxfun(@times,effect',Z);
    P = 1./(1+exp(-F));
    P(P < P_TRUNC) = P_TRUNC;
    P(P > 1 - P_TRUNC) = 1 - P_TRUNC;
    W = P.*(1-P);
    workRes = F + (Y - P) ./ W;

    elbo = -0.5*sum(sum(W.*(workRes-F).^2)) -0.5*sum(sum(effect_var.*(W'*(X.^2))));

    % divergence of prior effect size
    log_ratio = bsxfun(@plus, log(effect_var_conditioned), log(invtausq_effect));
    div = 1 - bsxfun(@times, effect_sq_conditioned, invtausq_effect);
    div = bsxfun(@plus, div, log_ratio);
    elbo = elbo + 0.5*sum(div*Z);

    % divergence of variable selection
    p = Z;
    p(p < P_TRUNC) = P_TRUNC;
    p(p > 1 - P_TRUNC) = 1 - P_TRUNC;
    elbo = elbo + prior_lodds'*Z - sum(log(1+exp(prior_lodds)));
    elbo = elbo - p'*log(p) - (1-p)'*log(1-p);

end
